import tensorflow as tf
from keras.datasets import mnist
import numpy as np

import matplotlib.pyplot as plt


class MNISTLoader:
    """
    加载 MNIST 数据集
    """

    def __init__(self):
        """
        构造器/构造方法
        """
        # 设置 keras 内部自带的 mnist 的简写
        # 加载数据集
        # 首次运行会 从 https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz 下载数据
        # 在 Window 系统保存到 C:\Users\当前登录的用户名\.keras\datasets
        # 在 Linux/Mac 系统保存到 /home/当前登录的用户名/.keras/datasets
        (self.train_data, self.train_label), (self.test_data, self.test_label) = mnist.load_data()

        # print(self.train_data[0])
        # 归一化处理
        # print(self.train_data[0]/255)
        # print(self.train_data.shape)

        # self.train_data = np.expand_dims(self.train_data.astype(np.float32) / 255.0, axis=-1)
        # self.test_data = np.expand_dims(self.test_data.astype(np.float32) / 255.0, axis=-1)
        # print(self.train_data.shape)

        # 为什么这里进行了类型转换？加快处理速度   C 自动类型转换 char short -> int
        self.train_label = self.train_label.astype(np.int32)
        self.test_label = self.test_label.astype(np.int32)
        # 训练数据的个数
        self.num_train_data = self.train_data.shape[0]
        self.num_test_data = self.test_data.shape[0]

        # for i in range(100):
        #     for j in range(10):
        #         print(self.train_label[10*i+j], end=" ")
        #     print("")

    def get_batch(self, batch_size):
        """
        随机获得指定个数的图形
        :param batch_size: 图片的个数
        :return:
        """
        index = np.random.randint(0, self.num_train_data, batch_size)
        return self.train_data[index, :], self.train_label[index]

    def get_train_data(self):
        """
        获取训练数据
        :return:
        """
        return self.train_data

    def get_train_label(self):
        """
        获取训练数据对应的标签
        :return:
        """
        return self.train_label

    def get_test_data(self):
        """
        获取测试数据
        :return:
        """
        return self.test_data

    def get_test_label(self):
        """
        获取测试数据对应的标签
        :return:
        """
        return self.test_label


if __name__ == '__main__':
    dataset = MNISTLoader()
    # (image, label) = dataset.get_batch(1)
    #
    # plt.axis("off")
    # plt.title(f"Label: {label[0]}")
    # plt.imshow(image[0], cmap="Greys")
    # plt.show()
    rows = 4
    columns = 5
    count = rows * columns
    (image, label) = dataset.get_batch(count)

    for i in range(count):
        plt.subplot(rows, columns, i + 1)
        plt.axis("off")
        plt.title(f"Label: {label[i]}")
        plt.imshow(image[i], cmap="Greys")
    plt.show()
